{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary modules\n",
    "# uncomment to get plots displayed in notebook\n",
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from classy import Class\n",
    "from scipy.optimize import fsolve\n",
    "from scipy.interpolate import interp1d\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# esthetic definitions for the plots\n",
    "font = {'size'   : 16, 'family':'STIXGeneral'}\n",
    "axislabelfontsize='large'\n",
    "matplotlib.rc('font', **font)\n",
    "matplotlib.mathtext.rcParams['legend.fontsize']='medium'\n",
    "plt.rcParams[\"figure.figsize\"] = [8.0,6.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#############################################\n",
    "#\n",
    "# User settings controlling the figure aspect\n",
    "#\n",
    "z_max_pk = 46000       # highest redshift involved\n",
    "k_per_decade = 400     # number of k values, controls final resolution\n",
    "k_min_tau0 = 40.       # this value controls the minimum k value in the figure (it is k_min * tau0)\n",
    "P_k_max_inv_Mpc =1.0   # this value is directly the maximum k value in the figure in Mpc\n",
    "tau_num_early = 2000   # number of conformal time values before recombination, controls final resolution\n",
    "tau_num_late = 200     # number of conformal time values after recombination, controls final resolution\n",
    "tau_ini = 10.          # first value of conformal time in Mpc\n",
    "tau_label_Hubble = 20. # value of time at which we want to place the label on Hubble crossing\n",
    "tau_label_ks = 40.     # value of time at which we want to place the label on sound horizon crossing\n",
    "tau_label_kd = 230.    # value of time at which we want to place the label on damping scale crossing\n",
    "#\n",
    "# Cosmological parameters and other CLASS parameters\n",
    "#\n",
    "common_settings = {# which output? transfer functions only\n",
    "                   'output':'mTk',\n",
    "                   # LambdaCDM parameters\n",
    "                   'h':0.67556,\n",
    "                   'omega_b':0.022032,\n",
    "                   'omega_cdm':0.12038,\n",
    "                   'A_s':2.215e-9,\n",
    "                   'n_s':0.9619,\n",
    "                   'tau_reio':0.0925,\n",
    "                   # Take fixed value for primordial Helium (instead of automatic BBN adjustment)\n",
    "                   'YHe':0.246,\n",
    "                   # other output and precision parameters\n",
    "                   'z_max_pk':z_max_pk,\n",
    "                   'k_per_decade_for_pk':k_per_decade,\n",
    "                   'k_per_decade_for_bao':k_per_decade,\n",
    "                   'k_min_tau0':k_min_tau0, # this value controls the minimum k value in the figure\n",
    "                   'perturbations_sampling_stepsize':'0.05',\n",
    "                   'P_k_max_1/Mpc':P_k_max_inv_Mpc,\n",
    "                   'compute damping scale':'yes', # needed to output and plot Silk damping scale\n",
    "                   'gauge':'newtonian'}\n",
    "\n",
    "###############\n",
    "#   \n",
    "# call CLASS \n",
    "#\n",
    "###############\n",
    "M = Class()\n",
    "M.set(common_settings)\n",
    "M.compute()\n",
    "#\n",
    "# define conformal time sampling array\n",
    "#\n",
    "times = M.get_current_derived_parameters(['tau_rec','conformal_age'])\n",
    "tau_rec=times['tau_rec']\n",
    "tau_0 = times['conformal_age']\n",
    "tau1 = np.logspace(math.log10(tau_ini),math.log10(tau_rec),tau_num_early)\n",
    "tau2 = np.logspace(math.log10(tau_rec),math.log10(tau_0),tau_num_late)[1:]\n",
    "tau2[-1] *= 0.999 # this tiny shift avoids interpolation errors\n",
    "tau = np.concatenate((tau1,tau2))\n",
    "tau_num = len(tau)\n",
    "#\n",
    "# use table of background and thermodynamics quantitites to define some functions \n",
    "# returning some characteristic scales\n",
    "# (of Hubble crossing, sound horizon crossing, etc.) at different time\n",
    "#\n",
    "background = M.get_background() # load background table\n",
    "#print background.viewkeys()\n",
    "thermodynamics = M.get_thermodynamics() # load thermodynamics table\n",
    "#print thermodynamics.viewkeys()    \n",
    "#\n",
    "background_tau = background['conf. time [Mpc]'] # read conformal times in background table\n",
    "background_z = background['z'] # read redshift\n",
    "background_aH = 2.*math.pi*background['H [1/Mpc]']/(1.+background['z'])/M.h() # read 2pi * aH in [h/Mpc]\n",
    "background_ks = 2.*math.pi/background['comov.snd.hrz.']/M.h() # read 2pi/(comoving sound horizon) in [h/Mpc]\n",
    "background_rho_m_over_r =\\\n",
    "    (background['(.)rho_b']+background['(.)rho_cdm'])\\\n",
    "    /(background['(.)rho_g']+background['(.)rho_ur']) # read rho_r / rho_m (to find time of equality)\n",
    "background_rho_l_over_m =\\\n",
    "    background['(.)rho_lambda']\\\n",
    "    /(background['(.)rho_b']+background['(.)rho_cdm']) # read rho_m / rho_lambda (to find time of equality)\n",
    "thermodynamics_tau = thermodynamics['conf. time [Mpc]'] # read confromal times in thermodynamics table\n",
    "thermodynamics_kd = 2.*math.pi/thermodynamics['r_d']/M.h() # read 2pi(comoving diffusion scale) in [h/Mpc]\n",
    "#\n",
    "# define a bunch of interpolation functions based on previous quantities\n",
    "#\n",
    "background_z_at_tau = interp1d(background_tau,background_z)\n",
    "background_aH_at_tau = interp1d(background_tau,background_aH)\n",
    "background_ks_at_tau = interp1d(background_tau,background_ks)\n",
    "background_tau_at_mr = interp1d(background_rho_m_over_r,background_tau)\n",
    "background_tau_at_lm = interp1d(background_rho_l_over_m,background_tau)\n",
    "thermodynamics_kd_at_tau = interp1d(thermodynamics_tau, thermodynamics_kd)\n",
    "#\n",
    "# infer arrays of characteristic quantitites calculated at values of conformal time in tau array\n",
    "#\n",
    "aH = background_aH_at_tau(tau) \n",
    "ks = background_ks_at_tau(tau)\n",
    "kd = thermodynamics_kd_at_tau(tau)\n",
    "#\n",
    "# infer times of R/M and M/Lambda equalities\n",
    "#\n",
    "tau_eq = background_tau_at_mr(1.)\n",
    "tau_lambda = background_tau_at_lm(1.)\n",
    "#\n",
    "# check and inform user whether intiial arbitrary choice of z_max_pk was OK\n",
    "max_z_needed = background_z_at_tau(tau[0])\n",
    "if max_z_needed > z_max_pk:\n",
    "    print ('you must increase the value of z_max_pk to at least ',max_z_needed)\n",
    "    () + 1  # this strange line is just a trick to stop the script execution there\n",
    "else:\n",
    "    print ('in a next run with the same values of tau, you may decrease z_max_pk from ',z_max_pk,' to ',max_z_needed)\n",
    "#\n",
    "# get transfer functions at each time and build arrays Theta0(tau,k) and phi(tau,k)\n",
    "#\n",
    "for i in range(tau_num):\n",
    "    one_time = M.get_transfer(background_z_at_tau(tau[i])) # transfer functions at each time tau\n",
    "    if i ==0:   # if this is the first time in the loop: create the arrays (k, Theta0, phi)\n",
    "        k = one_time['k (h/Mpc)']\n",
    "        k_num = len(k)\n",
    "        Theta0 = np.zeros((tau_num,k_num))\n",
    "        phi = np.zeros((tau_num,k_num))\n",
    "    Theta0[i,:] = 0.25*one_time['d_g'][:]\n",
    "    phi[i,:] = one_time['phi'][:]\n",
    "#\n",
    "# find the global extra of Theta0(tau,k) and phi(tau,k), used to define color code later\n",
    "#\n",
    "Theta_amp = max(Theta0.max(),-Theta0.min()) \n",
    "phi_amp = max(phi.max(),-phi.min()) \n",
    "#\n",
    "# reshaping of (k,tau) necessary to call the function 'pcolormesh'\n",
    "#\n",
    "K,T = np.meshgrid(k,tau)\n",
    "#\n",
    "# inform user of the size of the grids (related to the figure resolution)\n",
    "#\n",
    "print ('grid size:',len(k),len(tau),Theta0.shape)\n",
    "#\n",
    "#################\n",
    "#\n",
    "# start plotting\n",
    "#\n",
    "#################\n",
    "#\n",
    "fig = plt.figure(figsize=(18,8)) \n",
    "#\n",
    "# plot Theta0(k,tau)\n",
    "#\n",
    "ax_Theta = fig.add_subplot(121)\n",
    "print ('> Plotting Theta_0')\n",
    "fig_Theta = ax_Theta.pcolormesh(K,T,Theta0,cmap='coolwarm',vmin=-Theta_amp,vmax=Theta_amp,shading='auto')\n",
    "print ('> Done')\n",
    "#\n",
    "# plot lines (characteristic times and scales)\n",
    "#\n",
    "ax_Theta.axhline(y=tau_rec,color='k',linestyle='-')\n",
    "ax_Theta.axhline(y=tau_eq,color='k',linestyle='-')\n",
    "ax_Theta.axhline(y=tau_lambda,color='k',linestyle='-')\n",
    "ax_Theta.plot(aH,tau,'r-',linewidth=2)\n",
    "ax_Theta.plot(ks,tau,color='#FFFF33',linestyle='-',linewidth=2)\n",
    "ax_Theta.plot(kd,tau,'b-',linewidth=2)\n",
    "#\n",
    "# dealing with labels\n",
    "#\n",
    "ax_Theta.set_title(r'$\\Theta_0$')\n",
    "ax_Theta.text(1.5*k[0],0.9*tau_rec,r'$\\mathrm{rec.}$')\n",
    "ax_Theta.text(1.5*k[0],0.9*tau_eq,r'$\\mathrm{R/M} \\,\\, \\mathrm{eq.}$')\n",
    "ax_Theta.text(1.5*k[0],0.9*tau_lambda,r'$\\mathrm{M/L} \\,\\, \\mathrm{eq.}$')\n",
    "ax_Theta.annotate(r'$\\mathrm{Hubble} \\,\\, \\mathrm{cross.}$',\n",
    "                  xy=(background_aH_at_tau(tau_label_Hubble),tau_label_Hubble),\n",
    "                  xytext=(0.1*background_aH_at_tau(tau_label_Hubble),0.8*tau_label_Hubble),\n",
    "                  arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "ax_Theta.annotate(r'$\\mathrm{sound} \\,\\, \\mathrm{horizon} \\,\\, \\mathrm{cross.}$',\n",
    "                  xy=(background_ks_at_tau(tau_label_ks),tau_label_ks),\n",
    "                  xytext=(0.07*background_aH_at_tau(tau_label_ks),0.8*tau_label_ks),\n",
    "                  arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "ax_Theta.annotate(r'$\\mathrm{damping} \\,\\, \\mathrm{scale} \\,\\, \\mathrm{cross.}$',\n",
    "                  xy=(thermodynamics_kd_at_tau(tau_label_kd),tau_label_kd),\n",
    "                  xytext=(0.2*thermodynamics_kd_at_tau(tau_label_kd),2.0*tau_label_kd),\n",
    "                  arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "#\n",
    "# dealing with axes\n",
    "#\n",
    "ax_Theta.set_xlim(k[0],k[-1])\n",
    "ax_Theta.set_xscale('log')\n",
    "ax_Theta.set_yscale('log')\n",
    "ax_Theta.set_xlabel(r'$k  \\,\\,\\, \\mathrm{[h/Mpc]}$')\n",
    "ax_Theta.set_ylabel(r'$\\tau   \\,\\,\\, \\mathrm{[Mpc]}$')\n",
    "ax_Theta.invert_yaxis()\n",
    "#\n",
    "# color legend\n",
    "#\n",
    "fig.colorbar(fig_Theta)\n",
    "#\n",
    "# plot phi(k,tau)\n",
    "#\n",
    "ax_phi = fig.add_subplot(122)\n",
    "ax_phi.set_xlim(k[0],k[-1])\n",
    "#ax_phi.pcolor(K,T,phi,cmap='coolwarm')\n",
    "print ('> Plotting phi')\n",
    "fig_phi = ax_phi.pcolormesh(K,T,phi,cmap='coolwarm',vmin=-0.,vmax=phi_amp,shading='auto')\n",
    "print ('> Done')\n",
    "#\n",
    "# plot lines (characteristic times and scales)\n",
    "#\n",
    "ax_phi.axhline(y=tau_rec,color='k',linestyle='-')\n",
    "ax_phi.axhline(y=tau_eq,color='k',linestyle='-')\n",
    "ax_phi.axhline(y=tau_lambda,color='k',linestyle='-')\n",
    "ax_phi.plot(aH,tau,'r-',linewidth=2)\n",
    "ax_phi.plot(ks,tau,color='#FFFF33',linestyle='-',linewidth=2)\n",
    "#\n",
    "# dealing with labels\n",
    "#\n",
    "ax_phi.set_title(r'$\\phi$')\n",
    "ax_phi.text(1.5*k[0],0.9*tau_rec,r'$\\mathrm{rec.}$')\n",
    "ax_phi.text(1.5*k[0],0.9*tau_eq,r'$\\mathrm{R/M} \\,\\, \\mathrm{eq.}$')\n",
    "ax_phi.text(1.5*k[0],0.9*tau_lambda,r'$\\mathrm{M/L} \\,\\, \\mathrm{eq.}$')\n",
    "ax_phi.annotate(r'$\\mathrm{Hubble} \\,\\, \\mathrm{cross.}$',\n",
    "                  xy=(background_aH_at_tau(tau_label_Hubble),tau_label_Hubble),\n",
    "                  xytext=(0.1*background_aH_at_tau(tau_label_Hubble),0.8*tau_label_Hubble),\n",
    "                  arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "ax_phi.annotate(r'$\\mathrm{sound} \\,\\, \\mathrm{horizon} \\,\\, \\mathrm{cross.}$',\n",
    "                  xy=(background_ks_at_tau(tau_label_ks),tau_label_ks),\n",
    "                  xytext=(0.07*background_aH_at_tau(tau_label_ks),0.8*tau_label_ks),\n",
    "                  arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "#\n",
    "# dealing with axes\n",
    "#\n",
    "ax_phi.set_xscale('log')\n",
    "ax_phi.set_yscale('log')\n",
    "ax_phi.set_xlabel(r'$k \\,\\,\\, \\mathrm{[h/Mpc]}$')\n",
    "ax_phi.set_ylabel(r'$\\tau \\,\\,\\, \\mathrm{[Mpc]}$')\n",
    "ax_phi.invert_yaxis()\n",
    "#\n",
    "# color legend\n",
    "#\n",
    "fig.colorbar(fig_phi)\n",
    "#\n",
    "# produce the PDF\n",
    "#\n",
    "#plt.show()\n",
    "plt.savefig('many_times.png',dpi=300)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
